#include <iostream>
#include <string>
#include <cstring>
#include <thread>
#include <arpa/inet.h>
#include <unistd.h>
#define BUFFER_SIZE 1024
#define PORT 8080

//TCP服务器搭建流程
// 1.创建listen套接字
// 2.bind协议家族、IP、端口
// 3.将lisen套接字设置为listen状态
// 4.accept获取链接
// 5.利用获取上来的套接字进行通信
void tcpServer()
{
    int sockfd,new_sockfd;
    struct sockaddr_in servaddr;
    char buffer[BUFFER_SIZE]={0};
    socklen_t len=sizeof(servaddr);
    if((sockfd=socket(AF_INET,SOCK_STREAM,0))<0)
    {
        perror("TCP socket failed");
        return;
    }

    memset(&servaddr,0,len);

    servaddr.sin_family=AF_INET;
    servaddr.sin_addr.s_addr=INADDR_ANY;
    servaddr.sin_port=htons(PORT);
    if(bind(sockfd,(struct sockaddr*)&servaddr,len)<0)
    {
        perror("TCP bind failed");
        close(sockfd);
        return;
    }
    if(listen(sockfd,3)<0)
    {
        perror("TCP listen failed");
        close(sockfd);
        return;
    }
    if((new_sockfd=accept(sockfd,(struct sockaddr*)&servaddr,&len))<0)
    {
        perror("TCP accept failed");
        close(sockfd);
        return;
    }
    while(true)
    {
        memset(&buffer,0,BUFFER_SIZE);
        if(read(new_sockfd,buffer,BUFFER_SIZE)<=0) break;
        std::cout<<"Received: " <<buffer<<std::endl;
        write(new_sockfd,buffer,strlen(buffer));
    }
    close(new_sockfd);
    close(sockfd);
}

//TCP客户端搭建流程
// 1.创建普通套接字
// 2.connect连接服务器
// 3.进行通信
void tcpClient()
{
    int sockfd;
    struct sockaddr_in servaddr;
    char buffer[BUFFER_SIZE];
    socklen_t len=sizeof(servaddr);

    if((sockfd=socket(AF_INET,SOCK_STREAM,0))<0)
    {
        perror("TCP socket failed");
        return;
    }
    memset(&servaddr,0,len);
    servaddr.sin_family=AF_INET;
    servaddr.sin_port=htons(PORT);
    inet_pton(AF_INET,"127.0.0.1",&servaddr.sin_addr);

    if(connect(sockfd,(struct sockaddr*)&servaddr,len)<0)
    {
        perror("TCP connect failed");
        close(sockfd);
        return;
    }
    while(true)
    {
        std::string message;
        std::cout<<"Enter message: ";
        std::getline(std::cin,message);

        write(sockfd,message.c_str(),message.length());
        memset(&buffer,0,BUFFER_SIZE);
        if(read(sockfd,buffer,BUFFER_SIZE)<=0) break;
        std::cout<<"Echo: "<<buffer<<std::endl;
    }
    close(sockfd);
}

// UDP服务器搭建流程
// 1.创建套接字
// 2.bind协议家族、IP、端口
// 3.进行通信
void udpServer()
{
    int sockfd;
    struct sockaddr_in servaddr,cliaddr;
    char buffer[BUFFER_SIZE]={0};

    if((sockfd=socket(AF_INET,SOCK_DGRAM,0))<0)
    {
        perror("UDP socket failed");
        return;
    }
    memset(&servaddr,0,sizeof(servaddr));
    memset(&cliaddr,0,sizeof(cliaddr));

    servaddr.sin_family=AF_INET;
    servaddr.sin_addr.s_addr=INADDR_ANY;
    servaddr.sin_port=htons(PORT);

    if(bind(sockfd,(struct sockaddr*)&servaddr,sizeof(servaddr))<0)
    {
        perror("UDP bind failed");
        return;
    }

    while(true)
    {
        socklen_t len=sizeof(cliaddr);
        memset(&cliaddr,0,BUFFER_SIZE);

        if(recvfrom(sockfd,buffer,BUFFER_SIZE,0,(struct sockaddr*)&cliaddr,&len)<=0)
            break;
        std::cout<<"Received: " <<buffer<<std::endl;
        sendto(sockfd,buffer,strlen(buffer),0,(struct sockaddr *)&cliaddr, len);
    }       
    close(sockfd);
}

// UDP客户端搭建流程
// 1.创建套接字
// 2.进行通信
void udpClient()
{
    int sockfd;
    struct sockaddr_in servaddr;
    char buffer[BUFFER_SIZE]={0};

    if((sockfd=socket(AF_INET,SOCK_DGRAM,0))<0)
    {
        perror("UDP socket failed");
        return;
    }

    memset(&servaddr,0,sizeof(servaddr));
    servaddr.sin_family=AF_INET;
    servaddr.sin_port=htons(PORT);
    inet_pton(AF_INET,"127.0.0.1",&servaddr.sin_addr);

    while(true){
        std::string message;
        std::cout<<"Enter message: ";
        std::getline(std::cin,message);
        sendto(sockfd,message.c_str(),message.length(),0,(struct sockaddr*)&servaddr,sizeof(servaddr));
        memset(buffer,0,BUFFER_SIZE);
        socklen_t len=sizeof(servaddr);
        if(recvfrom(sockfd,buffer,BUFFER_SIZE,0,(struct sockaddr*)&servaddr,&len)<=0)
            break;
        std::cout<<"Echo: "<<buffer<<std::endl;
    }
    close(sockfd);
}


int main() {
    std::cout << "Choose mode: \n1. TCP Server \n2. TCP Client \n3. UDP Server \n4. UDP Client\n";
    int choice;
    std::cin >> choice;
    std::cin.ignore(); // Clear newline

    switch (choice) {
        case 1:
            tcpServer();
            break;
        case 2:
            tcpClient();
            break;
        case 3:
            udpServer();
            break;
        case 4:
            udpClient();
            break;
        default:
            std::cout << "Invalid choice." << std::endl;
            break;
    }

    return 0;
}